#include "StdAfx.h"
#include "AESEngine.h"


/***
 *     Constructors 
 */

AESEngine::AESEngine(void)
{
	setup();
	buildMathTables();
}


AESEngine::~AESEngine(void)
{
	removeMathTables();
}

/***
 *     Implements the IAESEngine interface 
 */

int AESEngine::
Encrypt(IAESKey* key, const byte in[], byte out[])
{
	int result = E_SUCCESS;

	if (!key || !in || !out){
		perror ("Encrypt invalid param");
		return E_INVALID_ARGS;
	}

	if (E_SUCCESS == 
			(result = ConfigureState(key)))
	{
		Encipher((byte*) in, out);
	}

	CloseState();

	return 0;
}


int AESEngine::
Decrypt(IAESKey* key, const byte in[], byte out[])
{
	if (!key || !in || !out){
		perror ("Decrypt invalid param");
		return E_INVALID_ARGS;
	}

	ConfigureState(key);

	// Decipher((byte*) in, out);

	CloseState();

	return 0;
}

/***
 *
 *  Helpers
 */

int AESEngine::
ConfigureState (IAESKey* key)
{
	int result = E_SUCCESS;

	// verify key is same as the configured key, if any 
	if (!verifyKeyIdentity(key)){
		key->PopulateKey();

		setKeyLength (key);
		setNumberRounds (key);

		// assign - key is reference counted?
		_key = key;

		if (E_SUCCESS == 
			(result = buildKeyTable()))
		{
			ExpandKey ();
		}
	}

	return result;
}

// builds the expanded key _mathTablesle that includes
// room for all the round keys
int AESEngine::
buildKeyTable(void)
{
	int result = E_SUCCESS;

	_w = new byte[AES_WORD_SIZE * Nb *(Nr+1)]; // room for the expanded key
	if (!_w) {
		perror ("buildKeyTable no memory error");
		return E_NO_MEMORY;
	}

	return result;
}

void AESEngine::
removeKeyTable(void)
{
	delete [] _w; 
	_w = NULL;
}


int AESEngine::
buildStateTable(void)
{
	int result = E_SUCCESS;

	_state = new byte*[AES_MAX_ROWS]; // the _state array
	if (!_state) {
		perror ("buildStateTable No memory error");
		return E_NO_MEMORY;
	}

	for (int i = 0; i < AES_MAX_ROWS; i++){
		_state[i] = new byte[Nb]; // the _state array
		_ASSERT(*_state);
	}
	
	return result;
}


void AESEngine::
removeStateTable(void)
{
	delete [] _state; 
	_state = NULL;
}

void AESEngine::
buildMathTables()
{
	_mathTables = new AESMath;
	_ASSERT(_mathTables);
}

void AESEngine::
removeMathTables()
{
	delete _mathTables;
}

/***
 *     Building blocks of the cipher operation 
 *
 *         - KeyExpansion	: Produces additional Keys used in the rounds of transformation.
 *		   - SubBytes		: Substitution increases the distance between the key and the cipher.
 *		   - ShiftRows		: adds to the diffusion of bits.
 *		   - MixColumns		: adds to the diffusion of bits.
 *		   - AddRoundKey	: xors the round key
 */


// Cipher: actual AES encrytion
void AESEngine::
Encipher(byte in[], byte out[]) 
{
	int result = E_SUCCESS;

	// byte counter for expanded key 
	_wCount = 0; 

	// creates a vector to hold the state
	if (E_SUCCESS != (result = buildStateTable ()))
	{
		return;
	}

	// sets up the plaintext for processing
	Copy(_state, in); 

	AddRoundKey(_state); 

	for (int round = 1; round < Nr; round++) 
	{
		printArray("Start round ", round , _state);

		SubBytes	(_state); // S-box substitution
		ShiftRows	(_state); // mix up rows
		MixColumns	(_state); // complicated mix of columns
		AddRoundKey	(_state); // xor with expanded key
	}
	
	printArray("Start round ", Nr, _state);

	SubBytes	(_state); // S-box substitution
	ShiftRows	(_state); // mix up rows
	AddRoundKey	(_state); // xor with expanded key
	Copy(out, _state);

	removeStateTable();
 }



// SubBytes: apply Sbox substitution to each byte of _state
void AESEngine::
SubBytes(byte** _state) 
{
	for (int row = 0; row < AES_MAX_ROWS; row++)
	{
		for (int col = 0; col < Nb; col++)
		{
			// substitution step using a table lookup
			_state[row][col] = 
				_mathTables->SBox(_state[row][col]);
		}
	}
}

// ShiftRows: simple circular shift of rows 1, 2, 3 by 1, 2, 3
void AESEngine::
ShiftRows(byte** _state) 
{
	byte* t = new byte[AES_WORD_SIZE];

	for (int r = 1; r < AES_MAX_ROWS; r++) 
	{
		for (int c = 0; c < Nb; c++)
		{
			t[c] = _state[r][(c + r)%Nb];
		}

		for (int c = 0; c < Nb; c++)
		{
			_state[r][c] = t[c];
		}
	}	
}

// AddRoundKey: xor a portion of expanded key with _state
void AESEngine::AddRoundKey(byte** _state) 
{
	for (int c = 0; c < Nb; c++)
	{
		for (int r = 0; r < AES_MAX_ROWS; r++)
		{
			_state[r][c] = (byte)(_state[r][c] ^ _w[_wCount++]);
		}
	}
}

// MixColumns: complex and sophisticated mixing of columns
void AESEngine::
MixColumns(byte** s) 
{
	int*	sp	= new int[4];
	byte	b02 = (byte)0x02, b03 = (byte)0x03;
	
	for (int c = 0; c < 4; c++) 
	{
		sp[0] = _mathTables->FFMul(b02,  s[0][c]) ^ _mathTables->FFMul(b03, s[1][c]) ^
								s[2][c] ^ s[3][c];

		sp[1] = 				s[0][c] ^ _mathTables->FFMul(b02, s[1][c]) ^
				_mathTables->FFMul(b03,  s[2][c]) ^ s[3][c];

		sp[2] = 				s[0][c] ^ s[1][c] ^
				_mathTables->FFMul(b02, 			s[2][c]) ^ _mathTables->FFMul(b03, s[3][c]);

		sp[3] = _mathTables->FFMul(b03,  s[0][c]) ^ s[1][c] ^
								s[2][c] ^ _mathTables->FFMul(b02, s[3][c]);

		for (int i = 0; i < 4; i++) {
			s[i][c] = (byte)(sp[i]);
		}
	}	
}

/***
 *		Helpers
 */

bool AESEngine::verifyKeyIdentity (IAESKey* key)
{
	if (!key) {
		perror ("verifyKeyIdentity invalid param");
		return 0;
	}

	return (_key == key);
};

void AESEngine::setKeyLength (IAESKey* key)
{
	// words in a key, = (4, or 6, or 8)
	Nk = key->getLengthKey ();
}

void AESEngine::setNumberRounds (IAESKey* key)
{
	// number of rounds corresponding to key length
	Nr = Nk + 6; 
}



int AESEngine::
CopyKey (void)
{
	int result = E_SUCCESS;

	if(!validateState())
	{
		perror ("CopyKey Invalid State");
		return E_FAILED;
	}

	// first just copy key to w
	int j = 0;

	while (j < AES_WORD_SIZE * Nk) {
		PKEY_STREAM aeskey = (PKEY_STREAM) ((AESKey*)_key)->getKeyStream();
		_w[j] = aeskey[j++];
	}

	return result;
}

int  AESEngine::
newAESWordBuffer(byte **outbuffer)
{
	int result = E_SUCCESS;
	byte* temp = new byte[4];

	if (!temp){
		perror ("newAESWordBuffer no memory");
		return E_NO_MEMORY;
	}

	*outbuffer = temp;

	return result;
}

// Expand key, byte-oriented code, but tracks words
int AESEngine::
ExpandKey (void)
{
	int result = E_SUCCESS;
	byte* temp = NULL;

	// get an AES word buffer
	if (E_SUCCESS != 
		(result = newAESWordBuffer(&temp))){
		return result; 
	}

	// first just copy key to w
	CopyKey ();
	printArray("key: ", _w, 16);
	
	// index positioned to create the expanded key set
	int j = AES_WORD_SIZE * Nk;

	int i;
	while(j < 4*Nb*(Nr+1)) 
	{
		i = j/4; // j is always multiple of 4 here

		// handle everything word-at-a time(AES_WORD_SIZE), 4 bytes at a time
		for (int iTemp = 0; iTemp < AES_WORD_SIZE; iTemp++)
		{
			temp[iTemp] = _w[j - AES_WORD_SIZE + iTemp];
		}

		// 
		if (i % Nk == 0) 
		{
			byte ttemp, tRcon;
			byte oldtemp0 = temp[0];

			for (int iTemp = 0; iTemp < AES_WORD_SIZE; iTemp++) 
			{
				if (iTemp == 3) {
					ttemp = oldtemp0;
				}
				else {
					ttemp = temp[iTemp+1];
				}

				if (iTemp == 0) {
					tRcon = _mathTables->Rcon(i/Nk);
				}
				else {
					tRcon = 0;
				}

				temp[iTemp] = (byte)(_mathTables->SBox(ttemp) ^ tRcon);
			}
		}
		else if (Nk > NWORDS_AES_KEY_196 && (i%Nk) == NWORDS_AES_KEY_128) 
		{
			for (int iTemp = 0; iTemp < AES_WORD_SIZE; iTemp++)
			{
				temp[iTemp] = _mathTables->SBox(temp[iTemp]);
			}
		}
	
		for (int iTemp = 0; iTemp < AES_WORD_SIZE; iTemp++)
		{
			_w[j+iTemp] = (byte)(_w[j - AES_WORD_SIZE * Nk + iTemp] ^ temp[iTemp]);
		}

		j = j + AES_WORD_SIZE;
	}

	return result;
}

bool AESEngine::
validateState()
{
	bool result = true;

	if (!_key || !_mathTables || (Nk < 4) || (Nk > 8)){
		perror ("validateState: E_INVALID_STATE");

		return false;
	}	

	return result;
}

void AESEngine::
CloseState()
{
	removeKeyTable();
}


// copy: copy in to _state
void AESEngine::
Copy(byte** _state, byte* in) 
{
	int inLoc = 0;
	for (int c = 0; c < Nb; c++){
		for (int r = 0; r < AES_WORD_SIZE; r++){
			_state[r][c] = in[inLoc++];
		}
	}
}

// copy: copy _state to out
void AESEngine::
Copy(byte* out, byte** _state) 
{
	int outLoc = 0;

	for (int c = 0; c < Nb; c++){
		for (int r = 0; r < AES_WORD_SIZE; r++){
			out[outLoc++] = _state[r][c];
		}
	}
}

// hex: print a byte as two hex digits
byte AESEngine::
Hex(byte a) 
{
	return dig[(a & 0xff) >> 4] + dig[a & 0x0f];
}



void  AESEngine::
printArray(char name[], byte a[], int length) 
{
	cout << name << " ";
	for (int i = 0; i < length; i++)
	{
		printf("%x ", a[i]);
	}

	cout << "\n";
}


void  AESEngine::
printArray(char name[], int round, byte** s) 
{
	cout << name << " " << round << ": ";
	for (int c = 0; c < Nb; c++){
		for (int r = 0; r < 4; r++){
			printf("%x ", s[r][c]);
		}
	}

	cout << "\n";
}

void AESEngine::
setup(void )
{
	Nk	= Nr = _wCount = 0;	 
	_w	=  NULL;
	_mathTables = NULL;
	_key	= NULL;
	_state = NULL;
};
